﻿import numpy as np
import tensorflow as tf
from tensorflow import keras

def data_set(stop):
  i=0
  while i<stop:
    x=np.random.randint(-1,2,size=(10,10))
    if x.sum()>0:
      y=1
    elif x.sum()<0:
      y=-1
    else:
      y=0
    i+=1
    x=x[:,:,np.newaxis]
    yield x,y

def model():
    inp=keras.layers.Input(shape=(10,10,1))
    x=keras.layers.Conv2D(100, 10)(inp)
    x=keras.layers.Flatten()(x)
    x=keras.layers.Dense(128)(x)
    outp=keras.layers.Dense(1,activation='tanh')(x)
    return keras.models.Model(inputs=inp, outputs=outp)

class example:
    def __init__(self):
        self.model=model()
        self.dset=tf.data.Dataset.from_generator(
                data_set,
                args=[1000],
                output_types=(tf.int32,tf.int32),
                output_shapes=([10,10,1], ())
                )
    def ex_compile(self):
        self.model.compile(
            optimizer=keras.optimizers.Adam(learning_rate=0.001),
            loss=keras.losses.MeanSquaredError(),
            metrics=['mse'])
    def ex_fit(self,batchSize,epochs):
        dst=self.dset.batch(batchSize,drop_remainder=True)
        self.model.fit(dst,epochs=epochs)
    def ex_predict(self,examples):
        return self.model.predict(examples)

test=example()
test.ex_compile()
test.ex_fit(100,200)

testSet=tf.data.Dataset.from_generator(
            data_set,
            args=[100],
            output_types=(tf.int32,tf.int32),
            output_shapes=([10,10,1], ())
            )

test_set=[]
test_y=[]
for x,y in testSet.take(50):
    test_set.append(x)
    test_y.append(y)
test_set=np.array(test_set)
test_y=np.array(test_y)

y_=test.ex_predict(test_set)
y_[y_<-0.3]=-1
y_[y_>0.3]=1
y_[abs(y_)!=1]=0
y_=np.array(y_,dtype=int).flatten()

print(y_==test_y)
